import os
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

tf.random.set_seed(22)
np.random.seed(22)
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'  显示所有的等级信息
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'  显示‘error’ 'warning' 信息
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'  显示‘error’信息
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
assert tf.__version__.startswith('2.')

# 构建数据集
# 常见单词数量
total_words = 10000
# 句子单词长度
max_review_len = 80
# 单词的表达维度
embedding_len = 100
batchsz = 128
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=total_words)
# x_train 和 x_test 都会被padding成为一个[b, 80]
x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=max_review_len)
x_test = keras.preprocessing.sequence.pad_sequences(x_test, maxlen=max_review_len)

db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
# 一个数据集最后的一个batch可能不够设定的batchsz，设置drop_remainder=True将最后一个batch丢弃掉
db_train = db_train.shuffle(1000).batch(batchsz, drop_remainder=True)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = db_test.batch(batchsz, drop_remainder=True)
print('x_train.shape: ', x_train.shape, tf.reduce_max(y_train), tf.reduce_min(y_train))


# 构建网络
class MyRNN(keras.Model):
    def __init__(self, units):
        super(MyRNN, self).__init__()

        # transform text to embedding representation
        # [b, 80] ==> [b, 80, 100]
        # 参数为：单词总数， 单词表达长度， 句子长度
        self.embedding = layers.Embedding(total_words, embedding_len,
                                          input_length=max_review_len)

        # [b, 80, 100] ==> [b, 80, h_dim] 这里h_dim用units表示
        # 语义提取
        self.rnn = keras.Sequential([
            # return_sequences=True返回上一层的状态， unroll=True ： 加速
            layers.LSTM(units, dropout=0.5, return_sequences=True, unroll=True),
            layers.LSTM(units, dropout=0.5, unroll=True)
        ])

        # fc [b, 64] => [b, 1]
        self.outlayer = layers.Dense(1)

    def call(self, inputs, training=None):
        """
        针对dropout:测试和训练网络并不相同
        net(x)和net(x, training=True) : training mode
        net(x, training=False) : test mode

        :param inputs: [b, 80]
        :param training:
        :return:
        """
        # [b, 80]
        x = inputs
        # embedding : [b, 80] => [b, 80, 100]
        x = self.embedding(x)
        # rnn cell compute
        # [b, 80, 100] => [b, 64]
        x = self.rnn(x)

        # out:[b, 64]
        x = self.outlayer(x)
        prob = tf.sigmoid(x)
        return prob


def main():
    units = 64
    epoch = 4
    model = MyRNN(units)
    model.compile(optimizer=keras.optimizers.Adam(0.001),
                  loss=tf.losses.BinaryCrossentropy(),
                  metrics=['accuracy'])
    model.fit(db_train, epochs=epoch, validation_data=db_test)
    model.evaluate(db_test)


if __name__ == '__main__':
    main()
